;;;; -*- Mode: Lisp; -*-


;;; ------------------------------------------------- 
;;; Matrix math.  Just enough to be dangerous, and almost all borrowed
;;; from GamePhysics.  Probably not quite generic enough for all types
;;; of matrix operations, but good enough to do cool game physics.  We
;;; hope.


;;; WARNING: Think carefully about this next statement, Mr. John.  If we want
;;; to pass external-arrays directly to GL then wouldn't gl-arrays need to be stored
;;; col-major?  Convince yourself one way or another and write it down somewhere!
;;;
;;; Note that external-arrays are stored by row-major index, but the matrices defined
;;; in this code are stored as column-major to better match with how gl does things.
;;;
;;;

(in-package :lib.gl)

(defparameter *default-matrix-constructor* 'array) ;; Either array or gl-array

(export '(m* m+ m-
	  get-column))
(defgeneric mref (m row &optional col))
(defgeneric row-major-mref (m idx))
(defgeneric m* (m1 m2 &rest args))
(defgeneric m+ (m1 m2 &rest args))
(defgeneric m- (m1 m2 &rest args))
(defgeneric vector-length (m))

;;;
;;; GL arrays are stored column major.
;;;

(defun column-vector-p (v)
  (typecase v
    (external-array (not (listp (dims v))))
    (vector t)
    (array (= (length (array-dimensions v)) 1))))

(defun print-std-matrix (matrix &optional (stream *standard-output*))
  (if (column-vector-p matrix)
      (iter (for x in-vector matrix)
	    (format stream "| ~9,2F |~%" x))
      (dotimes (r (matrix-rows matrix))
	(format stream "| ")
	(dotimes (c (matrix-cols matrix))
	  (format stream "~9,2F" (mref matrix r c)))
	(format stream " |")
	(terpri stream))))

(defmethod print-matrix ((m array) &rest args)
  (apply #'print-std-matrix m args))


;;;; 

(defgeneric matrix-rows (matrix))
(defgeneric matrix-cols (matrix))

(defmethod matrix-rows ((matrix array))
  (array-dimension matrix 0))

(defmethod matrix-cols ((matrix array))
  (array-dimension matrix 1))

(defmethod mref ((matrix array) row &optional col)
  (aref matrix row col))

(defmethod (setf mref) (val (matrix array) &rest indices)
  (setf (aref matrix (first indices) (second indices))
        val))

(defmethod row-major-mref ((matrix array) idx)
  (row-major-aref matrix idx))

(defmethod (setf row-major-mref) (val (matrix array) idx)
  (setf (row-major-aref matrix idx) val))

(iterate:defmacro-clause (FOR var IN-MATRIX matrix)
  (alexandria:with-unique-names (index m end)
    `(progn
       (with ,m = ,matrix)
       (with ,end = (* (matrix-rows ,matrix) (matrix-cols ,matrix)))
       (with ,index = -1)
       (for ,var next (progn (incf ,index)
			     (if (>= ,index ,end) (finish))
			     (row-major-mref ,m ,index))))))

(defun matrix-equal (m1 m2)
  (iter (for e1 in-matrix m1)
	(for e2 in-matrix m2)
	(always (< (- e1 e2) +epsilon+))))

;(defsetf mref set-mref)

(defun get-row (matrix rowidx)
  (apply #'make-vector 
	 (loop for i from 0 below (matrix-cols matrix)
	    collecting (mref matrix rowidx i))))

(defun set-row (matrix rowidx vals)
  (loop for j from 0 below (vector-length vals)
     doing (setf (mref matrix rowidx j) (mref vals j)))
  matrix)

(defun get-column (matrix colidx)
  (apply #'make-vector 
	 (loop for i from 0 below (matrix-rows matrix)
	    collecting (mref matrix i colidx))))

(defun set-column (matrix colidx vals)
  (loop for i from 0 below (vector-length vals)
        doing (setf (mref matrix i colidx) (mref vals i)))
  matrix)

(defmethod element-type ((m array))
  (type-of (aref m 0 0)))

(defmethod element-count ((m array))
  (* (matrix-rows m) (matrix-cols m)))

(defun submatrix (a i0 j0 i1 j1)
  "Return the subblock of matrix <a> between the indices (i0, j0) and (i1, j1)."
  (let ((b (make-matrix (1+ (- i1 i0)) (1+ (- j1 j0))
			:element-type (element-type a))))
    (iter (for i from i0 to i1)
	  (finally (return b))
	  (iter (for j from j0 to j1)
		(setf (mref b (- i i0) (- j j0)) (mref a i j))))))

(defun make-square-matrix (&rest objects)
  (let* ((d (truncate (sqrt (length objects))))
         (arr (make-matrix d d :element-type (type-of (first objects)))))
    (dotimes (i d arr)
      (dotimes (j d)
        (setf (mref arr i j)
              (nth (+ j (* i d)) objects))))))

(defun zero-matrix (rows &optional (cols nil) &key (element-type 'single-float))
  (let ((zero (coerce +zero+ element-type)))
    (if cols
	(make-matrix cols rows :initial-element zero)
	(make-matrix rows rows :initial-element zero))))

(defun zero-matrix-p (a)
  (iter (for x in-matrix a)
	(always (< (abs x) +epsilon+))))

(defun square-matrix-p (matrix)
  (= (matrix-rows matrix) (matrix-cols matrix)))

(defun identity-matrix (rows &key (element-type 'single-float))
  (let* ((one (coerce +one+ element-type))
	 (zero (coerce +zero+ element-type))
	 (matrix (make-matrix rows rows :initial-element zero)))
    (loop for i upto (1- rows)
       doing (setf (mref matrix i i) one))
    matrix))

(defun matrix3-by-cols (v1 v2 v3)
  (let ((newm (make-matrix 3 3 :initial-element +zero+)))
    (set-column newm 0 v1)
    (set-column newm 1 v2)
    (set-column newm 2 v3)
    newm))

(defun matrix3-by-rows (r1 r2 r3)
  (let ((newm (make-matrix 3 3 :initial-element +zero+)))
    (set-row newm 0 r1)
    (set-row newm 1 r2)
    (set-row newm 2 r3)
    newm))

(defun %matrix (objects)
  (assert (listp (car objects)))
  (assert (= (length (first objects))
	     (length (car (last objects)))))
  (let* ((rows (length objects))
	 (cols (length (first objects)))
	 (arr (make-matrix rows cols :element-type (type-of (caar objects))))
	 (i 0))
    (dolist (row objects arr)
      (set-row arr i (apply #'make-vector row));(map 'list #'(lambda (x) (coerce x 'double-float))
					      ; row)))
      (incf i))))

(defun copy-matrix (a)
  (let* ((ni (matrix-rows a))
	 (nj (matrix-cols a))
	 (b (make-matrix ni nj :element-type (element-type a))))
    (dotimes (i ni b)
      (dotimes (j nj b)
	(setf (mref b i j) (mref a i j))))))

(defun make-matrix (rows columns &key (initial-element +zero+) (element-type nil))
  (case *default-matrix-constructor*
    (array  (if element-type
		(make-array (list rows columns) :element-type element-type)
		(make-array (list rows columns) :initial-element initial-element)))
    (external-array (if element-type
		(make-external-array (list rows columns) :element-type element-type)
		(make-external-array (list rows columns) :initial-element initial-element)))))

;;;; Iterate Clauses ;;;;
(iterate:defmacro-clause (FOR var IN-MATRIX-COL m)
    (alexandria:with-unique-names (index matrix end)
      `(progn
	 (with ,matrix = ,m)
	 (with ,end = (matrix-cols ,matrix))
	 (with ,index = -1)
	 (for ,var next (progn (incf ,index)
			       (if (>= ,index ,end) (finish))
			       (get-column ,matrix ,index))))))

(iterate:defmacro-clause (FOR var IN-MATRIX-ROW m)
    (alexandria:with-unique-names (index matrix end)
      `(progn
	 (with ,matrix = ,m)
	 (with ,end = (matrix-rows ,matrix))
	 (with ,index = -1)
	 (for ,var next (progn (incf ,index)
			       (if (>= ,index ,end) (finish))
			       (get-row ,matrix ,index))))))

(defun random-square-matirx (rank &optional (scale 1.0))
  "Return matrix of rank <rank> filled with random values scaled by <scale>."
  (let ((m (make-matrix 12 12)))
    (dotimes (i rank m)
      (dotimes (j rank)
	(setf (mref m i j) (random scale))))))

;;;; Matrix division ;;;;

(defmethod m/ ((matrix array) (k number) &rest args)
  (declare (ignore args))
  (let ((r (matrix-rows matrix)))
    (if (not (= k 0))
        (let ((kinv (/ +one+ k)))
          (dotimes (i (* r r) matrix)
            (setf (row-major-aref matrix i)
                  (* (row-major-aref matrix i) kinv))))
	(dotimes (i (* r r) matrix)
	  (setf (row-major-aref matrix i) most-positive-double-float)))))

;;;; Matrix multiplication ;;;;

(defun %mult-matrices-general (m1 m2)
  (let* ((n (matrix-rows m1))
	 (m (matrix-cols m1))
	 (p (matrix-cols m2))
	 (q (make-matrix n p :initial-element (coerce +zero+ (element-type m1)))))
					;(persistent-zero-matrix n p (coerce +zero+ (element-type m1)))))
    (dotimes (i n q)
      (dotimes (j p)
	(let ((s 0))
	  (dotimes (k m)
	    (incf s (* (mref m1 i k) (mref m2 k j))))
	  (setf (mref q i j) s))))))

	    ;; (handler-case (incf (mref q i j) fac)
	    ;;   (ccl::array-element-type-error () (let ((etype (type-of (mref q i j))))
	    ;; 					  (incf (mref q i j) (coerce fac etype)))))))))))

;;; Only works with square matrices, yes?
(defun %mult-matrices (m p)
  (if (not (and (square-matrix-p m)
		(square-matrix-p p)))
      (%mult-matrices-general m p)
      (let* ((r (matrix-rows m))
	     (c (matrix-cols p))
	     ;(zero (coerce +zero+ (element-type m)))
	     (q (make-matrix r c :initial-element (coerce +zero+ (element-type m)))))
	(dotimes (row r q)
	  (dotimes (col c)
	    (dotimes (mid (matrix-cols m))
	      (incf (mref q row col)
		    (* (mref m row mid)
		       (mref p mid col)))))))))

(defun %mult-matvec (m v)
  ;(unless (= (vector-length v) (matrix-cols m))
  ;  (error "Number of rows, ~A, for vector ~A do not equal the number of columns, ~A, in matrix ~A."
  ;         (length v) v (matrix-rows m) m))
  (let* ((zero (coerce +zero+ (element-type m)))
	 (rv (make-array (vector-length v) :initial-element zero)))
    (dotimes (row (matrix-rows m) rv)
      (let ((sum 0))
	(dotimes (col (vector-length v))
	  (incf sum (* (mref m row col) (mref v col))))
        (setf (mref rv row) sum)))))
              

(defun %mult-vecmat (v m)
  (unless (and (= (vector-length v) (matrix-cols m))
	       (= (matrix-rows m) 1))
    (error "To multiply with column vector ~A, ~A needs to be a row vector with ~A columns."
	   v m (vector-length v)))
  (let* ((s (vector-length v))
	 (ra (make-matrix s s :initial-element +zero+)))
    (dotimes (i s ra)
      (dotimes (j s)
	(setf (mref ra i j) (* (mref v i) (mref m 0 j)))))))

(defun %mul-matrices-scalar (k p)
  (let* ((r (matrix-rows p))
	 (m (copy-matrix p)))
    (dotimes (i (* r r) m)
      (mult-assign (row-major-mref m i) k))))


(defmethod m* ((k number) (c number) &rest args)
  (if (not (null args))
      (apply #'m* (* k c) (first args) (rest args))
    (* k c)))

;; --- ;;;

(defmethod m* ((matrix array) (k number) &rest args)
  (apply #'m* k matrix args))

(defmethod m* ((k number) (matrix array) &rest args)
  (let ((r (%mul-matrices-scalar k matrix)))
    (if (not (null args))
        (apply #'m* r (first args) (rest args))
	r)))

(defmethod m* ((p array) (q array) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-matrices p q) (first rest) (rest rest))
    (%mult-matrices p q)))

(defmethod m* ((p array) (q vector) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-matvec p q) (first rest) (rest rest))
      (%mult-matvec p q)))

(defmethod m* ((q vector) (p array) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-vecmat p q) (first rest) (rest rest))
      (%mult-vecmat q p)))

;;;; ---------------------------

(defun %sub-mat (m1 m2)
  (unless (= (matrix-rows m1) (matrix-rows m2))
    (error "Cannot subtract matrices of different sizes."))
  (let* ((r (matrix-rows m1))
         (p (copy-matrix m1)))
    (dotimes (i (* r r) p)
      (setf (row-major-mref p i)
            (- (row-major-mref m1 i)
               (row-major-mref m2 i))))))

(defmethod m- ((m1 array) (m2 array) &rest args)
  (if (not (null args))
      (apply #'m- (%sub-mat m1 m2) (first args) (rest args))
    (%sub-mat m1 m2)))

(defun %add-mat (m1 m2)
  (unless (= (matrix-rows m1) (matrix-rows m2))
    (error "Cannot add matrices of different sizes."))
  (let* ((r (matrix-rows m1))
         (p (copy-matrix m1)))
    (dotimes (i (* r r) p)
      (setf (row-major-mref p i)
            (+ (row-major-mref m1 i)
               (row-major-mref m2 i))))))

(defmethod m+ ((m1 array) (m2 array) &rest args)
  (if (not (null args))
      (apply #'m+ (%add-mat m1 m2) (first args) (rest args))
    (%add-mat m1 m2)))

;;; Matrix transpose
(defun transpose (m)
  "Return the transpose of a matrix."
  (if (column-vector-p m)
      (let* ((c (vector-length m))
	     (mt (make-matrix 1 c :element-type (element-type m))))
	(iter (for x in-vector m)
	      (for i from 0)
	      (setf (mref mt 0 i) x)
	      (finally (return mt))))
      (let* ((r (matrix-rows m))
	     (c (matrix-cols m))
	     (mt (zero-matrix c r :element-type (element-type m))))
	(dotimes (row r mt)
	  (dotimes (col c)
	    (setf (mref mt col row) (mref m row col)))))))

;;;

;;;

(defun make-tensor-product (u v)
  (if (and (= (vector-length u) 3)
	   (= (vector-length v) 3))
      (make-square-matrix (* (mref u 0) (mref v 0))
			  (* (mref u 0) (mref v 1))
			  (* (mref u 0) (mref v 2))
			  (* (mref u 1) (mref v 0))
			  (* (mref u 1) (mref v 1))
			  (* (mref u 1) (mref v 2))
			  (* (mref u 2) (mref v 0))
			  (* (mref u 2) (mref v 0))
			  (* (mref u 2) (mref v 2)))
      (m* u (transpose v))))

(defun matrix-from-axis-angle (axis angle &optional (radians t))
  (let* ((fcos (cos (if radians angle (deg-to-rad angle))))
         (fsin (sin (if radians angle (deg-to-rad angle))))
         (oneminuscos (- 1 fcos))
         (x2 (* (x axis) (x axis)))
         (y2 (* (y axis) (y axis)))
         (z2 (* (z axis) (z axis)))
         (xym (* (x axis) (y axis) oneminuscos))
         (xzm (* (x axis) (z axis) oneminuscos))
         (yzm (* (y axis) (z axis) oneminuscos))
         (xsin (* (x axis) fsin))
         (ysin (* (y axis) fsin))
         (zsin (* (z axis) fsin)))
    (make-square-matrix (+ (* x2 oneminuscos) fcos)
                        (- xym zsin)
                        (+ xzm ysin)
                        (+ xym zsin)
                        (+ (* y2 oneminuscos) fcos)
                        (- yzm xsin)
                        (- xzm ysin)
                        (+ yzm xsin)
                        (+ (* z2 oneminuscos) fcos))))

;;;;

(defun test-make-matrix ()
  (%matrix '((1.0 2.0 3.0)
	     (4.0 5.0 6.0)
	     (7.0 8.0 9.0))))

(defun test-matrix-constructors ()
  (print-matrix (make-square-matrix 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0))
  (print-matrix (zero-matrix 5))
  (print-matrix (identity-matrix 8))
  (print-matrix (matrix3-by-rows #(1 2 3) #(4 5 6) #(7 8 9)))
  (print-matrix (matrix3-by-cols #(1 2 3) #(4 5 6) #(7 8 9)))
  (print-matrix (copy-matrix (make-square-matrix 1.0 2.0 3.0 4.0)))
  (assert (square-matrix-p (make-square-matrix 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0)))
  (assert (zero-matrix-p (zero-matrix 10))))

(defun test-matrix-access ()
  (let ((m (test-make-matrix)))
    (format t "~A~%" (mref m 1 1))
    (format t "~A~%" (get-row m 0))
    (format t "~A~%" (get-column m 2))
    (print-matrix (submatrix m 1 0 2 2))))

(defun test-general-mult ()
  (let ((e (%matrix '((1.0 0.0 0.0)
		      (-2.0 1.0 0.0)
		      (0.0 0.0 1.0))))
	(a (%matrix '((2.0 1.0 1.0)
		      (4.0 -6.0 0.0)
		      (-2.0 7.0 2.0)))))
    (%mult-matrices-general e a)))

(defun test-mat-vec-mult ()
  (let ((a (%matrix '((1.0 0.0 0.0)
		      (-2.0 1.0 0.0)
		      (0.0 0.0 1.0))))
	(v (make-vector 2.0 4.0 -2.0)))
    (%mult-matvec a v)))

(defun test-matrix-mult ()
  (let ((e (%matrix '((1.0 0.0 0.0)
		      (-2.0 1.0 0.0)
		      (0.0 0.0 1.0))))
	(a (%matrix '((2.0 1.0 1.0)
		      (4.0 -6.0 0.0)
		      (-2.0 7.0 2.0))))
	(v (make-vector 2.0 4.0 -2.0)))
    (print-matrix (m* e a))
    (terpri)
    (print-matrix (m* 2 e))
    (terpri)
    (print-matrix (m* a v))))

(defun test-matrix-add-sub ()
  (let ((e (%matrix '((1.0 0.0 0.0)
		      (-2.0 1.0 0.0)
		      (0.0 0.0 1.0))))
	(a (%matrix '((2.0 1.0 1.0)
		      (4.0 -6.0 0.0)
		      (-2.0 7.0 2.0)))))
    (print-matrix (m+ e a))
    (terpri)
    (print-matrix (m- e a))))

(defun test-transpose ()
  (let ((e (%matrix '((1.0 0.0 0.0)
		      (-2.0 1.0 0.0)
		      (0.0 0.0 1.0)))))
    (print-matrix (transpose e))))
	
(defun test-submatrix-1 ()
  (let ((a (test-make-matrix)))
    (assert (matrix-equal (%matrix '((1.0 2.0)
				     (4.0 5.0)))
			  (submatrix a 0 0 1 1)))
    (assert (matrix-equal (%matrix '((5.0 6.0)
				     (8.0 9.0)))
			  (submatrix a 1 1 2 2)))
    (assert (matrix-equal (%matrix '((4.0 5.0 6.0)
				     (7.0 8.0 9.0)))
			  (submatrix a 1 0 2 2)))))
